package com.chj.gaoji.class14;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.TreeSet;

public class Code06_CountofRangeSum {

	public int countRangeSum02(int[] nums, int lower, int upper) {
		if (nums == null || nums.length == 0) {
			return 0;
		}

		int numsLen = nums.length;
		int answer = 0;
		for (int i = 0; i < numsLen; ++i) {
			long sum = 0;
			for (int j = i; j < numsLen; ++j) {
				sum += nums[j];
				if (sum >= lower && sum <= upper)
					++answer;
			}
		}

		return answer;
	}

	public int countRangeSum01(int[] nums, int lower, int upper) {
		if (nums == null || nums.length == 0) {
			return 0;
		}
		List<Long> prefixSum = new ArrayList<>();
		prefixSum.add(0L);
		long sum = 0;
		int ans = 0;
		for (int i = 0; i < nums.length; i++) {
			sum += nums[i];
			int t1 = right_bound(prefixSum, sum - (long) lower);
			int t2 = left_bound(prefixSum, sum - (long) upper);
			ans += t1 - t2 + 1;
			int ta = left_bound(prefixSum, sum);
			prefixSum.add(ta, sum);
		}
		return ans;
	}

	public static int left_bound(List<Long> prefixSum, long l) {
		int left = 0, right = prefixSum.size() - 1;
		while (left <= right) {
			int mid = left + (right - left) / 2;
			if (prefixSum.get(mid) < l) {
				left = mid + 1;
			} else if (prefixSum.get(mid) > l) {
				right = mid - 1;
			} else if (prefixSum.get(mid) == l) {
				// 别返回，锁定左侧边界
				right = mid - 1;
			}
		}
		// // 最后要检查 left 越界的情况
		// if (left >= nums.length || nums[left] != target)
		// return -1;
		return left;
	}

	public static int right_bound(List<Long> prefixSum, long l) {
		int left = 0, right = prefixSum.size() - 1;
		while (left <= right) {
			int mid = left + (right - left) / 2;
			if (prefixSum.get(mid) < l) {
				left = mid + 1;
			} else if (prefixSum.get(mid) > l) {
				right = mid - 1;
			} else if (prefixSum.get(mid) == l) {
				// 别返回，锁定右侧边界
				left = mid + 1;
			}
		}
		// // 最后要检查 right 越界的情况
		// if (right < 0 || nums[right] != target)
		// return -1;
		return right;
	}

	public int countRangeSum(int[] nums, int lower, int upper) {
		// lower <= sum[j] - sum[i - 1] <= upper
		// sum[j] - upper <= sum[i - 1] <= sum[j] - lower
		int N = nums.length;

		long[] preNums = new long[N * 3 + 1];
		int prei = 0;
		preNums[prei++] = 0L;
		long sum = 0;
		for (int i = 0; i < N; i++) {
			sum += nums[i];
			preNums[prei++] = sum - upper;
			preNums[prei++] = sum - lower;
			preNums[prei++] = sum;
		}

		// SegmentTree seg = new SegmentTree(origin);
		int S = 1; // 整个区间的开始位置，规定从1开始，不从0开始 -> 固定
		// int SN = origin.length; // 整个区间的结束位置，规定能到N，不是N-1 -> 固定
		int root = 1; // 整棵树的头节点位置，规定是1，不是0 -> 固定
		// // 区间生成，必须在[S,N]整个范围上build
		// seg.build(S, SN, root);

		HashMap<Long, Integer> solArr = index(preNums);
		int SN = solArr.size();
		SegmentTree seg = new SegmentTree(SN);
		seg.build(S, SN, root);

		int ans = 0;
		sum = 0;
		int sni = 0;

		int pos0 = solArr.get(preNums[sni++]);
		seg.add(pos0, pos0, 1, S, SN, root);
		for (int i = 0; i < N; i++) {
			sum += nums[i];

			int ref1 = solArr.get(preNums[sni++]), ref2 = solArr.get(preNums[sni++]);
			ans += seg.query(ref1, ref2, S, SN, root);

			int pos1 = solArr.get(preNums[sni++]);
			seg.add(pos1, pos1, 1, S, SN, root);
		}
		return ans;
	}

	/**
	 * --------------------------------下面是离散化常规操作--------------------------------
	 */

	public HashMap<Long, Integer> index(long[] preNums) {
		TreeSet<Long> pos = new TreeSet<>();
		for (long posSub : preNums) {
			pos.add(posSub);
		}

		HashMap<Long, Integer> map = new HashMap<>();
		int count = 0;
		for (long index : pos) {
			map.put(index, ++count);
		}

		return map;
	}

	public class SegmentTree {
		// arr[]为原序列的信息从0开始，但在arr里是从1开始的
		// sum[]模拟线段树维护区间和
		// lazy[]为累加懒惰标记
		// change[]为更新的值
		// update[]为更新慵懒标记
		private int MAXN;
		private int[] cnt;

		public SegmentTree(int size) {
			MAXN = size + 1;
			cnt = new int[MAXN << 2];
		}

		// 在初始化阶段，先把sum数组，填好
		// 在arr[l~r]范围上，去build，1~N，
		// rt : 这个范围在sum中的下标
		public void build(int l, int r, int rt) {
			if (l == r) {
				cnt[rt] = 0;
				return;
			}
			int mid = (l + r) >> 1;
			build(l, mid, rt << 1);
			build(mid + 1, r, rt << 1 | 1);
			pushUp(rt);
		}

		private void pushUp(int rt) {
			cnt[rt] = cnt[rt << 1] + cnt[rt << 1 | 1];
		}

		public void update(int L, int R, int C, int l, int r, int rt) {
			if (L <= l && r <= R) {
				cnt[rt] = C;
				return;
			}
			// 当前任务躲不掉，无法懒更新，要往下发
			int mid = (l + r) >> 1;

			if (L <= mid) {
				update(L, R, C, l, mid, rt << 1);
			}
			if (R > mid) {
				update(L, R, C, mid + 1, r, rt << 1 | 1);
			}
			pushUp(rt);
		}

		// L..R -> 任务范围 ,所有的值累加上C
		// l,r -> 表达的范围
		// rt 去哪找l，r范围上的信息
		public void add(int L, int R, int C, int l, int r, int rt) {
			// 任务的范围彻底覆盖了，当前表达的范围
			if (L <= l && r <= R) {
				cnt[rt] += C;
				return;
			}
			// 要把任务往下发
			// 任务 L, R 没有把本身表达范围 l,r 彻底包住
			int mid = (l + r) >> 1;

			// 左孩子是否需要接到任务
			if (L <= mid) {
				add(L, R, C, l, mid, rt << 1);
			}
			// 右孩子是否需要接到任务
			if (R > mid) {
				add(L, R, C, mid + 1, r, rt << 1 | 1);
			}
			// 左右孩子做完任务后，我更新我的sum信息
			pushUp(rt);
		}

		public long query(int L, int R, int l, int r, int rt) {
			if (L <= l && r <= R) {
				return cnt[rt];
			}
			int mid = (l + r) >> 1;
			long ans = 0;
			if (L <= mid) {
				ans += query(L, R, l, mid, rt << 1);
			}
			if (R > mid) {
				ans += query(L, R, mid + 1, r, rt << 1 | 1);
			}
			return ans;
		}
	}

	public static void main(String[] args) {
		{
			int[] nums = { -2, 5, -1 };
			int lower = -2, upper = 2;

			Code06_CountofRangeSum code03_CountofRangeSum = new Code06_CountofRangeSum();

			System.out.println(code03_CountofRangeSum.countRangeSum(nums, lower, upper));
		}
	}
}
